from llvmlite import ir, binding

class CodeGenerator:
    def __init__(self):
        # 创建一个LLVM模块
        self.module = ir.Module(name='example_module')
        # 创建一个构建器
        self.builder = None
        # 函数列表
        self.functions = {}
        # 局部变量映射
        self.locals = {}
        # 全局变量映射
        self.globals = {}
        # 函数返回类型
        self.return_type = None
        # 存储函数参数
        self.params = {}
    
    def visit(self, node):
        """访问AST节点，根据节点类型调用相应的处理方法"""
        method_name = f'visit_{node.node_type}'
        visitor = getattr(self, method_name, self.generic_visit)
        return visitor(node)
    
    def generic_visit(self, node):
        """默认访问方法，当没有特定的处理方法时调用"""
        raise Exception(f'没有访问方法: {node.node_type}')
    
    def visit_Program(self, node):
        """访问程序节点"""
        # 重置所有状态（确保每次运行都从头开始）
        self.globals = {}
        # 不要重置self.functions，因为我们需要在处理函数调用时使用它
        
        # 首先分离函数定义和其他语句
        function_defs = []
        all_stmts = []
        
        print("开始处理程序节点，总语句数:", len(node.children))
        
        # 先检查整个AST结构，确保函数定义被正确识别
        for stmt in node.children:
            node_type = getattr(stmt, 'node_type', 'unknown')
            if node_type == 'Function':
                function_defs.append(stmt)
                func_name = stmt.value.get('name', 'unknown')
                print(f"  - 找到函数定义: {func_name}")
            else:
                all_stmts.append(stmt)
                if node_type == 'VarDeclaration':
                    var_name = stmt.value.get('name', 'unknown')
                    var_type = stmt.value.get('type', 'unknown')
                    print(f"  - 找到变量声明: {var_name}: {var_type}")
                else:
                    print(f"  - 找到其他语句: {node_type}")
        
        # 先处理所有函数定义
        print("\n处理函数定义...")
        for func_def in function_defs:
            self.visit(func_def)
        
        # 创建main函数作为程序入口
        print("\n创建main函数...")
        main_func_type = ir.FunctionType(ir.IntType(32), [])
        main_func = ir.Function(self.module, main_func_type, name='main')
        entry_block = main_func.append_basic_block(name='entry')
        self.builder = ir.IRBuilder(entry_block)
        
        # 重置局部变量和参数映射
        self.locals = {}
        self.params = {}
        
        # 为所有变量声明创建空间并注册到全局变量表
        print("\n第一阶段：为所有变量声明创建空间...")
        for stmt in all_stmts:
            if stmt.node_type == 'VarDeclaration':
                var_name = stmt.value['name']
                var_type_str = stmt.value['type']
                
                # 确定变量类型
                if var_type_str == 'int':
                    var_type = ir.IntType(32)
                elif var_type_str == 'float':
                    var_type = ir.FloatType()
                elif var_type_str == 'bool':
                    var_type = ir.IntType(1)
                
                # 创建变量空间
                var = self.builder.alloca(var_type, name=var_name)
                # 注册到全局变量表
                self.globals[var_name] = var
                print(f"  - 创建变量空间: {var_name}，已添加到全局变量表")
        
        # 打印当前全局变量表状态
        print("\n当前全局变量表:", list(self.globals.keys()))
        
        # 处理所有语句（变量初始化和其他语句）
        print("\n第二阶段：处理所有语句...")
        for stmt in all_stmts:
            node_type = getattr(stmt, 'node_type', 'unknown')
            if node_type == 'VarDeclaration':
                var_name = stmt.value['name']
                print(f"  - 处理变量初始化: {var_name}")
            self.visit(stmt)
        
        # 如果main函数没有明确的返回语句，添加默认返回0
        if not self.builder.block.is_terminated:
            self.builder.ret(ir.Constant(ir.IntType(32), 0))
        
        return self.module
    
    def visit_Function(self, node):
        """访问函数定义节点"""
        # 获取函数名和返回类型
        func_name = node.value['name']
        return_type_str = node.value['return_type']
        
        # 确定返回类型
        if return_type_str == 'int':
            return_type = ir.IntType(32)
        elif return_type_str == 'float':
            return_type = ir.FloatType()
        elif return_type_str == 'bool':
            return_type = ir.IntType(1)
        else:
            return_type = ir.VoidType()
        
        self.return_type = return_type
        
        # 获取参数类型
        params = node.children[0]
        arg_types = []
        for param in params:
            param_type_str = param.value['type']
            if param_type_str == 'int':
                arg_types.append(ir.IntType(32))
            elif param_type_str == 'float':
                arg_types.append(ir.FloatType())
            elif param_type_str == 'bool':
                arg_types.append(ir.IntType(1))
        
        # 创建函数类型
        func_ty = ir.FunctionType(return_type, arg_types)
        # 在模块中添加函数
        func = ir.Function(self.module, func_ty, name=func_name)
        # 保存函数引用
        self.functions[func_name] = func
        
        # 创建函数体基本块
        block = func.append_basic_block(name='entry')
        self.builder = ir.IRBuilder(block)
        
        # 保存参数和函数信息
        self.locals = {}
        self.params = {}
        for i, param in enumerate(params):
            param_name = param.value['name']
            self.params[param_name] = func.args[i]  # 参数直接保存为值类型
        
        # 访问函数体
        body = node.children[1]
        self.visit(body)
        
        # 如果函数没有明确的返回语句且返回类型不是void，添加默认返回
        if not self.builder.block.is_terminated and return_type != ir.VoidType():
            if return_type == ir.IntType(32):
                self.builder.ret(ir.Constant(ir.IntType(32), 0))
            elif return_type == ir.FloatType():
                self.builder.ret(ir.Constant(ir.FloatType(), 0.0))
            elif return_type == ir.IntType(1):
                self.builder.ret(ir.Constant(ir.IntType(1), 0))
        
        return func
    
    def visit_Block(self, node):
        """访问代码块节点"""
        # 保存当前局部变量映射的副本
        old_locals = self.locals.copy()
        
        # 访问代码块中的所有语句
        for stmt in node.children:
            self.visit(stmt)
        
        # 恢复局部变量映射
        self.locals = old_locals
        
        return None
    
    def visit_VarDeclaration(self, node):
        """访问变量声明节点"""
        var_name = node.value['name']
        var_type_str = node.value['type']
        
        # 确定变量类型
        if var_type_str == 'int':
            var_type = ir.IntType(32)
            default_value = ir.Constant(ir.IntType(32), 0)
        elif var_type_str == 'float':
            var_type = ir.FloatType()
            default_value = ir.Constant(ir.FloatType(), 0.0)
        elif var_type_str == 'bool':
            var_type = ir.IntType(1)
            default_value = ir.Constant(ir.IntType(1), 0)
        
        # 判断是否在main函数中（全局作用域）
        is_global_scope = self.builder and self.builder.block.function.name == 'main'
        
        # 检查变量是否已经在符号表中
        if var_name in self.globals:
            # 全局变量已经在visit_Program中创建
            var = self.globals[var_name]
            print(f"  - 变量 '{var_name}' 已存在于全局变量表中")
        elif var_name in self.locals:
            # 局部变量已经存在
            var = self.locals[var_name]
            print(f"  - 变量 '{var_name}' 已存在于局部变量表中")
        else:
            # 创建变量空间
            var = self.builder.alloca(var_type, name=var_name)
            # 根据作用域注册到相应的符号表
            if is_global_scope:
                self.globals[var_name] = var
                print(f"  - 变量 '{var_name}' 作为全局变量创建")
            else:
                self.locals[var_name] = var
                print(f"  - 变量 '{var_name}' 作为局部变量创建")
        
        # 处理初始化值
        if node.children:
            print(f"  - 处理变量 '{var_name}' 的初始化值")
            # 有初始化值
            value = self.visit(node.children[0])
            # 类型转换
            if isinstance(value.type, ir.FloatType) and isinstance(var_type, ir.IntType):
                value = self.builder.fptosi(value, var_type)
            elif isinstance(value.type, ir.IntType) and isinstance(var_type, ir.FloatType):
                value = self.builder.sitofp(value, var_type)
            self.builder.store(value, var)
        else:
            # 无初始化值，使用默认值
            self.builder.store(default_value, var)
        
        return var
    
    def visit_Assignment(self, node):
        """访问赋值语句节点"""
        target = node.children[0]
        value = self.visit(node.children[1])
        
        # 获取目标变量
        target_name = target.value['name']
        
        # 先检查局部变量
        if target_name in self.locals:
            target_var = self.locals[target_name]
        # 再检查全局变量
        elif target_name in self.globals:
            target_var = self.globals[target_name]
        else:
            raise Exception(f'未定义的变量: {target_name}')
        
        # 获取目标变量类型
        target_type = target_var.type.pointee
        
        # 类型转换
        if isinstance(value.type, ir.FloatType) and isinstance(target_type, ir.IntType):
            value = self.builder.fptosi(value, target_type)
        elif isinstance(value.type, ir.IntType) and isinstance(target_type, ir.FloatType):
            value = self.builder.sitofp(value, target_type)
        
        # 执行赋值
        self.builder.store(value, target_var)
        
        return value
    
    def visit_BinaryOp(self, node):
        """访问二元运算节点"""
        op = node.value['op']
        left = self.visit(node.children[0])
        right = self.visit(node.children[1])
        
        # 确保类型一致
        if isinstance(left.type, ir.FloatType) and isinstance(right.type, ir.IntType):
            right = self.builder.sitofp(right, left.type)
        elif isinstance(left.type, ir.IntType) and isinstance(right.type, ir.FloatType):
            left = self.builder.sitofp(left, right.type)
        
        # 根据运算符生成相应的LLVM指令
        if op == '+':
            if isinstance(left.type, ir.FloatType):
                return self.builder.fadd(left, right)
            else:
                return self.builder.add(left, right)
        elif op == '-':
            if isinstance(left.type, ir.FloatType):
                return self.builder.fsub(left, right)
            else:
                return self.builder.sub(left, right)
        elif op == '*':
            if isinstance(left.type, ir.FloatType):
                return self.builder.fmul(left, right)
            else:
                return self.builder.mul(left, right)
        elif op == '/':
            if isinstance(left.type, ir.FloatType):
                return self.builder.fdiv(left, right)
            else:
                return self.builder.sdiv(left, right)
        elif op == '%':
            # 确保操作数都是整数类型
            if isinstance(left.type, ir.IntType) and isinstance(right.type, ir.IntType):
                # 如果右操作数是布尔类型(i1)，转换为i32
                if right.type.width == 1:
                    right = self.builder.zext(right, ir.IntType(32))
                # 确保左操作数也是i32
                if left.type.width != 32:
                    left = self.builder.sext(left, ir.IntType(32))
                return self.builder.srem(left, right)
            else:
                raise Exception('模运算只能用于整数类型')
        elif op == '==':
            if isinstance(left.type, ir.FloatType):
                return self.builder.fcmp_ordered('==', left, right)
            else:
                return self.builder.icmp_signed('==', left, right)
        elif op == '!=':
            if isinstance(left.type, ir.FloatType):
                return self.builder.fcmp_ordered('!=', left, right)
            else:
                return self.builder.icmp_signed('!=', left, right)
        elif op == '<':
            if isinstance(left.type, ir.FloatType):
                return self.builder.fcmp_ordered('<', left, right)
            else:
                return self.builder.icmp_signed('<', left, right)
        elif op == '<=':
            if isinstance(left.type, ir.FloatType):
                return self.builder.fcmp_ordered('<=', left, right)
            else:
                return self.builder.icmp_signed('<=', left, right)
        elif op == '>':
            if isinstance(left.type, ir.FloatType):
                return self.builder.fcmp_ordered('>', left, right)
            else:
                return self.builder.icmp_signed('>', left, right)
        elif op == '>=':
            if isinstance(left.type, ir.FloatType):
                return self.builder.fcmp_ordered('>=', left, right)
            else:
                return self.builder.icmp_signed('>=', left, right)
        elif op == '&&':
            return self.builder.and_(left, right)
        elif op == '||':
            return self.builder.or_(left, right)
        
        raise Exception(f'不支持的运算符: {op}')
    
    def visit_UnaryOp(self, node):
        """访问一元运算节点"""
        op = node.value['op']
        expr = self.visit(node.children[0])
        
        if op == '!':
            return self.builder.not_(expr)
        elif op == '-':
            if isinstance(expr.type, ir.FloatType):
                return self.builder.fsub(ir.Constant(expr.type, 0.0), expr)
            else:
                return self.builder.sub(ir.Constant(expr.type, 0), expr)
        
        raise Exception(f'不支持的一元运算符: {op}')
    
    def visit_Identifier(self, node):
        """访问标识符节点"""
        name = node.value['name']
        
        # 首先检查是否是参数
        if hasattr(self, 'params') and name in self.params:
            # 参数直接返回，不需要加载
            return self.params[name]
        
        # 然后检查是否是局部变量
        if name in self.locals:
            # 加载局部变量值
            return self.builder.load(self.locals[name], name)
        
        # 最后检查是否是全局变量
        if name in self.globals:
            # 加载全局变量值
            return self.builder.load(self.globals[name], name)
        
        # 调试信息
        print(f"调试: 找不到变量 '{name}'，当前作用域:")
        print(f"  - 参数: {list(self.params.keys())}")
        print(f"  - 局部变量: {list(self.locals.keys())}")
        print(f"  - 全局变量: {list(self.globals.keys())}")
        
        raise Exception(f'未定义的变量: {name}')
    
    def visit_NumberLiteral(self, node):
        """访问数字字面量节点"""
        value = node.value['value']
        return ir.Constant(ir.IntType(32), value)
    
    def visit_FloatLiteral(self, node):
        """访问浮点数字面量节点"""
        value = node.value['value']
        return ir.Constant(ir.FloatType(), value)
    
    def visit_BooleanLiteral(self, node):
        """访问布尔字面量节点"""
        value = node.value['value']
        return ir.Constant(ir.IntType(1), 1 if value else 0)
    
    def visit_FunctionCall(self, node):
        """访问函数调用节点"""
        func_name = node.value['name']
        
        # 手动处理参数，确保能正确访问全局变量
        args = []
        for arg in node.children:
            # 先保存当前的locals和params，以便在处理参数时能正确查找全局变量
            old_locals = self.locals.copy()
            old_params = self.params.copy()
            
            try:
                # 处理参数
                arg_value = self.visit(arg)
                args.append(arg_value)
            finally:
                # 恢复locals和params
                self.locals = old_locals
                self.params = old_params
        
        # 查找函数
        if func_name not in self.functions:
            # 检查是否是外部函数（如print）
            if func_name == 'print':
                # 创建print函数声明
                if args and isinstance(args[0].type, ir.IntType):
                    if args[0].type.width == 1:
                        # 布尔值
                        print_ty = ir.FunctionType(ir.VoidType(), [ir.IntType(8).as_pointer()])
                        print_fn = ir.Function(self.module, print_ty, name='printf')
                        format_str = ir.Constant(ir.ArrayType(ir.IntType(8), 2), bytearray(b"%d\0"))
                    else:
                        # 整数
                        print_ty = ir.FunctionType(ir.VoidType(), [ir.IntType(8).as_pointer()])
                        print_fn = ir.Function(self.module, print_ty, name='printf')
                        format_str = ir.Constant(ir.ArrayType(ir.IntType(8), 2), bytearray(b"%d\0"))
                else:
                    # 浮点数
                    print_ty = ir.FunctionType(ir.VoidType(), [ir.IntType(8).as_pointer()])
                    print_fn = ir.Function(self.module, print_ty, name='printf')
                    format_str = ir.Constant(ir.ArrayType(ir.IntType(8), 4), bytearray(b"%f\0"))
                
                # 全局字符串
                format_global = ir.GlobalVariable(self.module, format_str.type, name='.str')
                format_global.global_constant = True
                format_global.initializer = format_str
                
                # 获取字符串指针
                format_ptr = self.builder.bitcast(format_global, ir.IntType(8).as_pointer())
                
                # 调用printf
                return self.builder.call(print_fn, [format_ptr] + args)
            
            raise Exception(f'未定义的函数: {func_name}')
        
        func = self.functions[func_name]
        return self.builder.call(func, args)
    
    def visit_IfStatement(self, node):
        """访问if语句节点"""
        print("开始处理if语句")
        # 生成条件表达式
        print("  - 处理条件表达式")
        condition = self.visit(node.children[0])
        print(f"  - 条件表达式生成完成，类型: {type(condition)}")
        
        # 获取当前函数
        func = self.builder.function
        print(f"  - 当前函数: {func.name}")
        
        # 创建基本块
        print("  - 创建基本块")
        then_block = func.append_basic_block(name='then')
        else_block = func.append_basic_block(name='else')
        merge_block = func.append_basic_block(name='ifcont')
        
        # 条件分支
        print("  - 生成条件分支")
        self.builder.cbranch(condition, then_block, else_block)
        
        # 访问then块
        print("  - 处理then块")
        self.builder.position_at_start(then_block)
        print(f"  - then块子节点数量: {len(node.children[1].children) if hasattr(node.children[1], 'children') else 0}")
        self.visit(node.children[1])
        if not self.builder.block.is_terminated:
            print("  - then块未终止，添加分支到merge")
            self.builder.branch(merge_block)
        
        # 访问else块（如果存在）
        print("  - 处理else块")
        if len(node.children) > 2:
            self.builder.position_at_start(else_block)
            self.visit(node.children[2])
            if not self.builder.block.is_terminated:
                self.builder.branch(merge_block)
        else:
            # 如果没有else块，直接跳转到merge块
            self.builder.position_at_start(else_block)
            self.builder.branch(merge_block)
        
        print("  - 定位到merge块")
        self.builder.position_at_start(merge_block)
        print("if语句处理完成")
    
    def visit_ForStatement(self, node):
        """访问for循环节点"""
        # 获取当前函数
        func = self.builder.function
        
        # 创建基本块
        preheader_block = self.builder.block
        loop_block = func.append_basic_block(name='loop')
        body_block = func.append_basic_block(name='loopbody')
        exit_block = func.append_basic_block(name='loopexit')
        
        # 处理初始化语句（preheader块）
        self.visit(node.children[0])
        self.builder.branch(loop_block)
        
        # 处理条件判断（loop块）
        self.builder.position_at_start(loop_block)
        condition = self.visit(node.children[1])
        self.builder.cbranch(condition, body_block, exit_block)
        
        # 处理循环体（body块）
        self.builder.position_at_start(body_block)
        self.visit(node.children[3])
        if not self.builder.block.is_terminated:
            # 处理更新语句
            self.visit(node.children[2])
            self.builder.branch(loop_block)
        
        # 设置出口块
        self.builder.position_at_start(exit_block)
        
        return None
    
    def visit_WhileStatement(self, node):
        """访问while语句节点"""
        # 获取当前函数
        func = self.builder.function
        
        # 创建基本块
        loop_entry = func.append_basic_block(name='loop_entry')
        loop_body = func.append_basic_block(name='loop_body')
        loop_exit = func.append_basic_block(name='loop_exit')
        
        # 跳转到循环入口
        self.builder.branch(loop_entry)
        
        # 循环入口：计算条件
        self.builder.position_at_start(loop_entry)
        condition = self.visit(node.children[0])
        self.builder.cbranch(condition, loop_body, loop_exit)
        
        # 循环体
        self.builder.position_at_start(loop_body)
        self.visit(node.children[1])
        if not self.builder.block.is_terminated:
            self.builder.branch(loop_entry)
        
        # 循环出口
        self.builder.position_at_start(loop_exit)
        
        return None
    
    def visit_ReturnStatement(self, node):
        """访问return语句节点"""
        if node.children:
            # 有返回值
            value = self.visit(node.children[0])
            # 类型转换
            if isinstance(value.type, ir.FloatType) and isinstance(self.return_type, ir.IntType):
                value = self.builder.fptosi(value, self.return_type)
            elif isinstance(value.type, ir.IntType) and isinstance(self.return_type, ir.FloatType):
                value = self.builder.sitofp(value, self.return_type)
            self.builder.ret(value)
        else:
            # 无返回值
            self.builder.ret_void()
        
        return None
    
    def generate_ir(self, ast):
        """生成LLVM IR"""
        return self.visit(ast)

def generate_llvm_ir(ast):
    """将AST转换为LLVM IR的入口函数"""
    codegen = CodeGenerator()
    return codegen.generate_ir(ast)